#!/bin/bash
#===============================================================================
# OpenFOAM并行强扩展性测试自动化脚本
# 功能：编译求解器，运行不同核数的并行测试，并统计计算时间
#
# 使用方法：
#   ./run_advection_scaling_test.sh [网格总数]
#
# 参数：
#   网格总数 - 可选，默认为10000000
#              脚本会自动计算sqrt(网格总数)并更新blockMeshDict
#
# 示例：
#   ./run_advection_scaling_test.sh              # 使用默认10000000个网格
#   ./run_advection_scaling_test.sh 1000000      # 使用1000000个网格 (1000x1000x1)
#   ./run_advection_scaling_test.sh 4            # 使用4个网格 (2x2x1)
#===============================================================================

# source $HOME/data/jiangxu/etc/bashrc_of2406


# 设置颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 获取脚本所在目录的绝对路径
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SOLVER_DIR="${SCRIPT_DIR}/advectionEqn/solver"
CASE_DIR="${SCRIPT_DIR}/advectionEqn/case"
LIB_DIR="${SCRIPT_DIR}/amgclFoam"

# 定义测试的核数数组
CORE_COUNTS=(9 18 36 72 144 288 576)
# CORE_COUNTS=(1 2 4)

# 定义网格总数（可以通过命令行参数传入，默认为10000000）
TOTAL_CELLS=${1:-10000000}

# 结果输出文件
# RESULT_FILE="${SCRIPT_DIR}/scaling_results_$(date +%Y%m%d_%H%M%S).txt"
RESULT_FILE="${SCRIPT_DIR}/scaling_results_advection.txt"

#===============================================================================
# 函数定义
#===============================================================================

# 打印带颜色的消息
print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 编译求解器
compile_solver() {
    print_info "开始编译求解器 advectionFoam..."
    cd "${SOLVER_DIR}" || exit 1
    
    # 清理之前的编译
    wclean
    
    # 编译求解器
    wmake 2>&1 | tee compile.log
    
    if [ $? -eq 0 ]; then
        print_success "求解器编译成功！"
        return 0
    else
        print_error "求解器编译失败！请查看 ${SOLVER_DIR}/compile.log"
        exit 1
    fi
}

# 编译库
compile_lib() {
    print_info "开始编译库 libamgclFoam..."
    cd "${LIB_DIR}" || exit 1
    
    # 清理之前的编译
    wclean
    
    # 编译库
    wmake 2>&1 | tee compile.log
    
    if [ $? -eq 0 ]; then
        print_success "amgclFoam库编译成功！"
        return 0
    else
        print_error "amgclFoam库编译失败！请查看 ${LIB_DIR}/compile.log"
        exit 1
    fi
}

# 修改decomposeParDict中的核数
modify_decompose_dict() {
    local nproc=$1
    local decompose_file="${CASE_DIR}/system/decomposeParDict"
    
    print_info "修改decomposeParDict为 ${nproc} 个子域..."
    
    # 备份原文件
    cp "${decompose_file}" "${decompose_file}.bak"
    
    # 修改numberOfSubdomains
    sed -i "s/^numberOfSubdomains.*/numberOfSubdomains ${nproc};/" "${decompose_file}"
    
    # 根据核数计算合适的分解方式（如果使用simple方法）
    # 这里可以根据需要调整分解策略
    # if [ $nproc -eq 9 ]; then
    #     sed -i "s/n.*( .* .* .* );/n               ( 3 3 1 );/" "${decompose_file}"
    # elif [ $nproc -eq 18 ]; then
    #     sed -i "s/n.*( .* .* .* );/n               ( 6 3 1 );/" "${decompose_file}"
    # elif [ $nproc -eq 36 ]; then
    #     sed -i "s/n.*( .* .* .* );/n               ( 6 6 1 );/" "${decompose_file}"
    # elif [ $nproc -eq 72 ]; then
    #     sed -i "s/n.*( .* .* .* );/n               ( 12 6 1 );/" "${decompose_file}"
    # elif [ $nproc -eq 144 ]; then
    #     sed -i "s/n.*( .* .* .* );/n               ( 12 12 1 );/" "${decompose_file}"
    # elif [ $nproc -eq 288 ]; then
    #     sed -i "s/n.*( .* .* .* );/n               ( 24 12 1 );/" "${decompose_file}"
    # elif [ $nproc -eq 576 ]; then
    #     sed -i "s/n.*( .* .* .* );/n               ( 24 24 1 );/" "${decompose_file}"
    # fi
}

# 修改blockMeshDict中的网格数
modify_mesh_size() {
    local total_cells=$1
    local blockmesh_file="${CASE_DIR}/system/blockMeshDict"
    
    print_info "计算网格划分参数（总网格数: ${total_cells}）..."
    
    # 计算平方根并四舍五入到最近的整数
    local mesh_per_side=$(echo "sqrt(${total_cells})" | bc)
    
    # 如果bc没有，使用awk
    if [ -z "$mesh_per_side" ] || [ "$mesh_per_side" == "0" ]; then
        mesh_per_side=$(awk -v n="${total_cells}" 'BEGIN{printf "%.0f", sqrt(n)}')
    fi
    
    print_info "每边网格数: ${mesh_per_side} x ${mesh_per_side} x 1"
    print_info "实际网格总数: $((mesh_per_side * mesh_per_side))"
    
    # 备份原文件（如果还没有备份的话）
    if [ ! -f "${blockmesh_file}.orig" ]; then
        cp "${blockmesh_file}" "${blockmesh_file}.orig"
    fi
    
    # 使用sed修改网格划分
    # 匹配形如 hex (...) (数字 数字 数字) 的行
    sed -i "s/\(hex ([^)]*) (\)[0-9]\+ [0-9]\+ [0-9]\+\()\)/\1${mesh_per_side} ${mesh_per_side} 1\2/" "${blockmesh_file}"
    
    print_success "blockMeshDict已更新为: (${mesh_per_side} ${mesh_per_side} 1)"
}

# 清理算例
clean_case() {
    print_info "清理算例目录..."
    cd "${CASE_DIR}" || exit 1
    
    # 清理processor目录
    rm -rf processor*
    
    # 清理时间步目录（保留0和constant）
    foamListTimes -rm
    
    # 清理日志文件
    rm -f log.*
    
    # 恢复初始条件
    rm -rf 0
    cp -r 0.orig 0
}

# 生成网格
generate_mesh() {
    print_info "使用blockMesh生成网格..."
    cd "${CASE_DIR}" || exit 1
    
    blockMesh > log.blockMesh 2>&1
    
    if [ $? -eq 0 ]; then
        print_success "网格生成成功！"
        return 0
    else
        print_error "网格生成失败！请查看 ${CASE_DIR}/log.blockMesh"
        exit 1
    fi
}

# 运行并行测试（通过Slurm提交）
run_parallel_test() {
    local nproc=$1
    local flag_file="${CASE_DIR}/.job_${nproc}_complete"
    local slurm_script="${SCRIPT_DIR}/slurm/job_adv_${nproc}.slurm"
    
    print_info "=========================================="
    print_info "准备提交 ${nproc} 核 Slurm 作业"
    print_info "=========================================="
    
    cd "${CASE_DIR}" || exit 1
    
    # 删除可能存在的旧标志文件
    rm -f "$flag_file"
    
    # 生成针对当前核数的 Slurm 脚本
    print_info "生成 Slurm 作业脚本: ${slurm_script}"
    
    cat > "$slurm_script" <<EOF
#!/bin/bash

#SBATCH --job-name=adv_${nproc}
#SBATCH --partition=intel
#SBATCH -n ${nproc}
#SBATCH --output=slog/adv_${nproc}.out
#SBATCH --error=slog/adv_${nproc}.err

ulimit -s unlimited
ulimit -l unlimited

module purge
source \$HOME/data/jiangxu/etc/bashrc_of2406
export LD_PRELOAD=\$HOME/data/jiangxu/cases/HUAWEI/mpiTest/mpiTimer/libmpitimer.so
export MPITIMER_ENABLE_BACKTRACE=1
export MPITIMER_SAMPLE_RATE=800
# 切换到算例目录
cd ${CASE_DIR}

# 只有多核时才需要区域分解
if [ ${nproc} -gt 1 ]; then
    echo "执行区域分解..."
    decomposePar > log.decomposePar 2>&1
    
    if [ \$? -ne 0 ]; then
        echo "区域分解失败！"
        touch ${flag_file}
        exit 1
    fi
fi

# 记录开始时间
start_time=\$(date +%s)

# 运行计算
if [ ${nproc} -eq 1 ]; then
    echo "开始串行计算（1核）..."
    advectionFoam > log.advectionFoam_${nproc} 2>&1
else
    echo "开始并行计算（${nproc}核）..."
    mpirun -np ${nproc} -x PATH -x LD_LIBRARY_PATH -x MPI_BUFFER_SIZE --bind-to core --map-by numa:PE=1 advectionFoam -parallel > log.advectionFoam_${nproc} 2>&1
fi

# 记录结束时间
end_time=\$(date +%s)
wall_time=\$((end_time - start_time))

echo "计算完成！墙钟时间: \${wall_time} s"

# 创建完成标志文件
touch ${flag_file}
EOF
    
    chmod +x "$slurm_script"
    
    # 提交 Slurm 作业
    print_info "提交 Slurm 作业..."
    local job_id=$(sbatch "$slurm_script" 2>&1 | grep -oP 'Submitted batch job \K\d+')
    
    if [ -n "$job_id" ]; then
        print_success "作业已提交，Job ID: ${job_id}"
    else
        print_warning "作业已提交（无法获取 Job ID）"
    fi
    
    # 等待作业完成（通过检测标志文件）
    print_info "等待作业完成（每10秒检测一次）..."
    local wait_count=0
    while [ ! -f "$flag_file" ]; do
        sleep 10
        wait_count=$((wait_count + 1))
        if [ $((wait_count % 6)) -eq 0 ]; then
            print_info "已等待 $((wait_count * 10)) 秒..."
        fi
    done
    
    print_success "${nproc}核作业已完成！"
    
    # 删除标志文件
    rm -f "$flag_file"
    
    # 提取ExecutionTime（CPU时间）和ClockTime（墙钟时间）
    cd "${CASE_DIR}"
    local exec_time=$(grep "ExecutionTime" log.advectionFoam_${nproc} | tail -1 | awk '{print $3}')
    local clock_time=$(grep "ClockTime" log.advectionFoam_${nproc} | tail -1 | awk '{print $7}')
    
    # 从Slurm输出文件中提取墙钟时间（如果有的话）
    local slurm_out="${SCRIPT_DIR}/slog/adv_${nproc}.out"
    local wall_time=""
    if [ -f "$slurm_out" ]; then
        wall_time=$(grep "墙钟时间:" "$slurm_out" | awk '{print $3}')
    fi
    
    # 输出结果
    echo "========================================" | tee -a "${RESULT_FILE}"
    echo "核数: ${nproc}" | tee -a "${RESULT_FILE}"
    if [ -n "$wall_time" ]; then
        echo "墙钟时间（Slurm统计）: ${wall_time} s" | tee -a "${RESULT_FILE}"
    fi
    echo "ExecutionTime (CPU): ${exec_time} s" | tee -a "${RESULT_FILE}"
    echo "ClockTime (求解器): ${clock_time} s" | tee -a "${RESULT_FILE}"
    echo "========================================" | tee -a "${RESULT_FILE}"
    echo "" | tee -a "${RESULT_FILE}"
    
    return 0
}

#===============================================================================
# 生成格式化的表格
#===============================================================================
generate_table() {
    local result_file=$1
    local table_file="${SCRIPT_DIR}/scaling_results_advection.txt"
    
    print_info "生成格式化表格..."
    
    # 解析结果文件，提取核数和时间
    local cores=()
    local times=()
    
    while IFS= read -r line; do
        if [[ $line =~ ^核数:[[:space:]]*([0-9]+) ]]; then
            cores+=("${BASH_REMATCH[1]}")
        elif [[ $line =~ ClockTime.*:[[:space:]]*([0-9.]+)[[:space:]]*s ]]; then
            times+=("${BASH_REMATCH[1]}")
        fi
    done < "$result_file"
    
    # 如果没有提取到ClockTime，尝试提取ExecutionTime
    if [ ${#times[@]} -eq 0 ]; then
        while IFS= read -r line; do
            if [[ $line =~ ExecutionTime.*:[[:space:]]*([0-9.]+)[[:space:]]*s ]]; then
                times+=("${BASH_REMATCH[1]}")
            fi
        done < "$result_file"
    fi
    
    # 计算加速比
    local base_time=${times[0]}
    local base_cores=${cores[0]}
    
    # 生成表格
    {
        echo "================================================================================"
        echo "OpenFOAM并行强扩展性测试结果表格"
        echo "================================================================================"
        echo "测试时间: $(date '+%Y-%m-%d %H:%M:%S')"
        echo "网格总数: ${TOTAL_CELLS}"
        echo ""
        printf "%-12s %-20s %-15s %-15s\n" "核心数" "计算用时 [s]" "加速比" "理想加速比"
        echo "--------------------------------------------------------------------------------"
        
        for i in "${!cores[@]}"; do
            local core=${cores[$i]}
            local time=${times[$i]}
            
            # 计算加速比 = 基准时间 / 当前时间
            local speedup=$(awk -v base="$base_time" -v curr="$time" 'BEGIN{printf "%.2f", base/curr}')
            
            # 理想加速比 = 核数比
            local ideal=$(awk -v base="$base_cores" -v curr="$core" 'BEGIN{printf "%.0f", curr/base}')
            
            printf "%-12s %-20s %-15s %-15s\n" "$core" "$time" "$speedup" "$ideal"
        done
        
        echo "================================================================================"
        echo ""
        echo "说明："
        echo "  - 加速比 = 基准时间 / 当前时间"
        echo "  - 理想加速比 = 当前核数 / 基准核数"
        echo "  - 并行效率 = 加速比 / 理想加速比 × 100%"
        echo ""
    } > "$table_file"
    
    # 同时输出到屏幕
    cat "$table_file"
    
    print_success "表格已保存到: ${table_file}"
}

#===============================================================================
# 主程序
#===============================================================================

main() {
    print_info "=========================================="
    print_info "OpenFOAM并行强扩展性测试"
    print_info "=========================================="
    echo ""
    
    # 显示测试配置
    print_info "测试配置:"
    print_info "  - 网格总数: ${TOTAL_CELLS}"
    print_info "  - 测试核数: ${CORE_COUNTS[@]}"
    echo ""
    
    # 初始化结果文件
    echo "OpenFOAM并行强扩展性测试结果" > "${RESULT_FILE}"
    echo "测试时间: $(date)" >> "${RESULT_FILE}"
    echo "算例目录: ${CASE_DIR}" >> "${RESULT_FILE}"
    echo "求解器: advectionFoam" >> "${RESULT_FILE}"
    echo "网格总数: ${TOTAL_CELLS}" >> "${RESULT_FILE}"
    echo "" >> "${RESULT_FILE}"
    
    # 步骤1：编译求解器
    compile_solver
    echo ""

    compile_lib
    echo ""
    
    # 步骤2：修改网格划分
    modify_mesh_size ${TOTAL_CELLS}
    echo ""
    
    # 步骤3：清理并生成网格（只需要一次）
    clean_case
    generate_mesh
    echo ""
    
    # 步骤4：循环测试不同核数
    for nproc in "${CORE_COUNTS[@]}"; do
        # 清理上次的processor目录
        rm -rf ${CASE_DIR}/processor*
        rm -rf ${CASE_DIR}/0
        cp -r ${CASE_DIR}/0.orig ${CASE_DIR}/0
        foamListTimes -rm > /dev/null 2>&1
        
        # 修改decomposeParDict
        modify_decompose_dict ${nproc}
        
        # 运行测试
        run_parallel_test ${nproc}
        
        echo ""
    done
    
    # 输出汇总
    print_success "=========================================="
    print_success "所有测试完成！"
    print_success "=========================================="
    print_info "结果已保存到: ${RESULT_FILE}"
    echo ""
    
    # 生成格式化表格
    generate_table "${RESULT_FILE}"
    echo ""
    
    cd "${CASE_DIR}" || exit 1
    
    # 清理processor目录
    rm -rf processor*
    
    # 清理生成的临时Slurm脚本（可选）
    print_info "清理临时Slurm脚本..."
    rm -f ${SCRIPT_DIR}/slurm/job_adv_*.slurm
    
    # 恢复decomposeParDict
    # if [ -f "${CASE_DIR}/system/decomposeParDict.bak" ]; then
    #     cp "${CASE_DIR}/system/decomposeParDict.bak" "${CASE_DIR}/system/decomposeParDict"
    #     rm "${CASE_DIR}/system/decomposeParDict.bak"
    # fi
    
    # 恢复blockMeshDict（可选，如果不想保留修改的话）
    # if [ -f "${CASE_DIR}/system/blockMeshDict.orig" ]; then
    #     cp "${CASE_DIR}/system/blockMeshDict.orig" "${CASE_DIR}/system/blockMeshDict"
    # fi
    
    print_success "所有清理工作完成！"
}

# 运行主程序
main

